package org.example.aigc.service.impl;

import com.alibaba.fastjson.JSONObject;
import org.example.aigc.service.ModelService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.env.Environment;
import org.springframework.core.io.FileSystemResource;
import org.springframework.http.*;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestTemplate;

import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;

@Service
public class ModelServiceImpl implements ModelService {

    /**
     * 更改语言模型的 API
     */
    private final String url = "http://172.26.94.12";

    private final RestTemplate restTemplate;

    private final String audioUploadDir;

    private final String videoUploadDir;

    @Autowired
    public ModelServiceImpl(
            RestTemplate restTemplate,
            Environment environment
    ) {
        this.restTemplate = restTemplate;
        this.audioUploadDir = Objects.requireNonNull(environment.getProperty("resources.audio-locations")).split(":", 2)[1];
        this.videoUploadDir = Objects.requireNonNull(environment.getProperty("resources.video-locations")).split(":", 2)[1];
    }

    public void setModel(String referWav, String referText) {
        String URI = url + ":9880/change_refer";
        // 向后端发送请求
        JSONObject request = new JSONObject();

        // 请求体
        referWav = "output/" + referWav;
        request.put("refer_wav_path", referWav);
        request.put("prompt_text", referText);
        request.put("prompt_language", "zh");
        // 请求变量
        JSONObject result = restTemplate.postForObject(URI, request, JSONObject.class);
        System.out.println("向服务器发送请求更改语音模型：" + result);
    }

    @Async
    public CompletableFuture<String> getWav(String content) {
        String URI = url + ":9880/";
        // 向后端发送请求
        JSONObject request = new JSONObject();

        // 请求体
        request.put("text", content);
        request.put("text_language", "zh");
        // 请求变量
        byte[] response = restTemplate.postForObject(URI, request, byte[].class);

        String fileName = UUID.randomUUID() + ".wav";

        try (FileOutputStream fos = new FileOutputStream(audioUploadDir + fileName)) {
            assert response != null;
            fos.write(response);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return CompletableFuture.completedFuture(fileName);
    }

    public String getMP4(String appearanceName, String wavName) {
        String URI = url + ":7863/execute_pipeline";

        // 创建请求体
        appearanceName = "in/" + appearanceName;
        wavName = "in/" + wavName;
        JSONObject body = getJsonObject(appearanceName, wavName);

        // 发送POST请求
        byte[] response = restTemplate.postForObject(URI, body, byte[].class);

        String fileName = UUID.randomUUID() + ".mp4";

        try (FileOutputStream fos = new FileOutputStream(videoUploadDir + fileName)) {
            assert response != null;
            fos.write(response);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return fileName;
    }

    public void uploadFile(String fileName, int isWav, int isGenWav) {
        Path filePath = Paths.get((isWav == 1 ? audioUploadDir : videoUploadDir) + fileName);
        FileSystemResource resource = new FileSystemResource(filePath);

        // 创建请求体
        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        body.add("file", resource);

        // 创建请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.MULTIPART_FORM_DATA);

        // 创建请求实体
        HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);
        ResponseEntity<String> responseEntity = restTemplate.exchange(
                isGenWav == 1 ? url + ":9880/upload" : url + ":7863/upload",
                HttpMethod.POST,
                requestEntity,
                String.class
        );
        System.out.println(responseEntity.getBody());
    }

    private static JSONObject getJsonObject(String appearanceName, String wavName) {
        JSONObject body = new JSONObject();
        body.put("video_file", appearanceName);
        body.put("vocal_file", wavName);
        body.put("quality", "Enhanced");
        body.put("output_height", "full resolution");
        body.put("wav2lip_version", "Wav2Lip");
        body.put("use_previous_tracking_data", "False");
        body.put("nosmooth", "False");
        body.put("U", 0);
        body.put("D", 10);
        body.put("L", 20);
        body.put("R", 0);
        body.put("size", 2.5);
        body.put("feathering", 5);
        body.put("mouth_tracking", "True");
        body.put("debug_mask", "False");
        body.put("batch_process", "False");
        return body;
    }
}
